from qutip import *
import numpy as np
sys.path.append('D:\Dropbox\Dokumente\PI3\Code\Python\quantum-tools')
import qutip_enhanced

qte = qutip_enhanced.QutipEnhanced()

def floor_odd(n):
    #0->0, 1->1, 2->1, 3->3, 4->3, ...
    return max(0, 2*np.floor((n+1)/2.) - 1)

def delta_alpha(hf, tau):
    return 2*np.pi*hf*tau

def alpha(n, hf, tau):
    return floor_odd(n)*(delta_alpha(hf, tau) + np.pi)

bloch = Bloch3d()
rho0_e = ket2dm(basis(2, 0))
rho0_n = ket2dm(basis(2, 0))
rho = tensor(rho0_e, rho0_n)

tau = 120
hf = 0.0123

def h(n):
    a = 2*np.pi*hf/2.
    b = np.exp(alpha(n, hf, tau)*1j)*0.5*2*np.pi/(2.*tau)
    h = Qobj([[a, 0, 0, 0],
             [0,  -a, 0, 0],
             [0,  0, 0, b.conj()],
             [0,  0, b, 0]])
    h.dims = [[2, 2], [2, 2]]
    return h

rho = tensor(rho0_e, rho0_n)
rho = qte.rotate(rho, rotation_axis={'y': -1}, angle=np.pi/2., rotated_spin=0)
rho = mesolve(h(0), rho, np.linspace(0, tau, 2), [], []).states[-1]
rho = qte.rotate(rho, rotation_axis={'y': -1}, angle=np.pi, rotated_spin=0)
rho = mesolve(h(1), rho, np.linspace(0, tau, 2), [], []).states[-1]
bloch.add_states(rho.ptrace(0))
bloch.show()